# Copyright (C) 2021.Huawei Technologies Co., Ltd. All rights reserved.
import os
import argparse
import ast
import subprocess

import moxing as mox

import security_utils

parser = argparse.ArgumentParser(description='Find Yunbao')

parser.add_argument('--data_url',
                    metavar='DIR',
                    default='/cache/data_url',
                    help='path to dataset')
parser.add_argument('--train_url',
                    default="/mindspore-dataset/output/",
                    type=str,
                    help="setting dir of training output")

parser.add_argument('--data_encrypted', type=ast.literal_eval, default=False,
                    help="data is encrypted, "
                         "need to be decrypted before training")

parser.add_argument('--code_encrypted', type=ast.literal_eval, default=False,
                    help="code is encrypted, "
                         "need to be decrypted before training")

parser.add_argument('--encrypt_model', type=ast.literal_eval, default=False,
                    help="model file need to be encrypted after training")

parser.add_argument('--user_dir_prefix', type=str, default='/cache',
                    choices=("/cache", "/home/ma-user/work/cache"),
                    help="user directory prefix")

args_opt = parser.parse_args()

CUR_RANK_ID = os.environ.get('RANK_ID')
USER_DIR_PREFIX = args_opt.user_dir_prefix
CACHE_TRAINING_URL = os.path.join(USER_DIR_PREFIX, f"training/{CUR_RANK_ID}/")

if not os.path.isdir(CACHE_TRAINING_URL):
    os.makedirs(CACHE_TRAINING_URL)

if __name__ == '__main__':
    real_path = os.path.join(USER_DIR_PREFIX, f'code/0')
    dew_config_file_name = 'dew.ini'
    if int(CUR_RANK_ID) == 0:
        if not os.path.exists(real_path):
            os.makedirs(real_path, 0o755)
        mox.file.copy(os.path.join(args_opt.data_url, dew_config_file_name),
                      os.path.join(real_path, dew_config_file_name))
        print(f"[{CUR_RANK_ID}] training data finish copy to {real_path}.")

    real_data_parent_path = real_path
    dew_config_file_path = os.path.join(real_data_parent_path,
                                        dew_config_file_name)
    cur_dir = os.path.dirname(os.path.realpath(__file__))
    cipher_code_dir = os.path.join(cur_dir, 'MNIST')
    if args_opt.code_encrypted:
        print(f'[{CUR_RANK_ID}] begin to get plain code...')
        plain_code_dir = security_utils.decrypt_dir(dew_config_file_path,
                                                    cipher_code_dir,
                                                    USER_DIR_PREFIX,
                                                    'decrypt_code.flag')
        print(f'[{CUR_RANK_ID}] get plain code successfully')
    else:
        plain_code_dir = cipher_code_dir
    real_path = plain_code_dir
    try:
        args_cmd = ['python3.7',
                    f'{real_path}/train.py'
                    ]
        for k, v in vars(args_opt).items():
            if k == 'code_encrypted' or not v:
                continue
            args_cmd.append(f'--{k}={v}')
        sub = subprocess.Popen(args_cmd, shell=False)
        sub.communicate()
    finally:
        if args_opt.code_encrypted:
            security_utils.delete_plain_dir(real_path)
